#include "aes.h"
#include <string.h>

/*
 * 嵌入式的最小 AES 实现（参考 tiny-AES-c），提供 AES-128 的密钥扩展、ECB、CBC 加密/解密工具函数。
 * 该实现适用于教学和小项目；生产环境请优先使用成熟的加密库（如 OpenSSL、libsodium）。
 */

// AES 的 S 盒，用于 SubBytes 步骤（字节替换表）
static const uint8_t sbox[256] = {
    /* 256-byte sbox table */
    0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76,
    0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0,
    0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15,
    0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75,
    0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84,
    0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf,
    0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8,
    0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2,
    0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73,
    0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb,
    0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79,
    0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08,
    0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a,
    0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e,
    0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf,
    0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16
};

// 轮常量，用于密钥扩展（KeyExpansion）
static const uint8_t Rcon[11] = {0x00,0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80,0x1B,0x36};

// KeyExpansion：把 16 字节的 AES 密钥扩展为 176 字节的轮密钥（用于 10 轮 AES-128）
static void KeyExpansion(uint8_t* RoundKey, const uint8_t* Key) {
    unsigned i, j;
    uint8_t tempa[4];

    // 第一个轮密钥就是原始密钥
    for (i = 0; i < 16; ++i) {
        RoundKey[i] = Key[i];
    }

    // 依次生成后续轮密钥
    for (i = 16; i < 176; i += 4) {
        for (j = 0; j < 4; ++j) {
            tempa[j] = RoundKey[i - 4 + j];
        }

        if (i % 16 == 0) {
            // RotWord：循环左移 1 字节
            uint8_t t = tempa[0];
            tempa[0] = tempa[1];
            tempa[1] = tempa[2];
            tempa[2] = tempa[3];
            tempa[3] = t;

            // SubWord：对每字节做 S-box 替换
            tempa[0] = sbox[tempa[0]];
            tempa[1] = sbox[tempa[1]];
            tempa[2] = sbox[tempa[2]];
            tempa[3] = sbox[tempa[3]];

            // 与 Rcon 异或
            tempa[0] = tempa[0] ^ Rcon[i/16];
        }

        for (j = 0; j < 4; ++j) {
            RoundKey[i + j] = RoundKey[i - 16 + j] ^ tempa[j];
        }
    }
}

// 以下是 AES 内部操作的辅助函数：AddRoundKey、SubBytes、ShiftRows、MixColumns 等
static void AddRoundKey(uint8_t state[4][4], const uint8_t* RoundKey) {
    // 将轮密钥与状态矩阵逐字节异或
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            state[j][i] ^= RoundKey[i * 4 + j];
        }
    }
}

// SubBytes：使用 S-box 对状态字节进行字节替换
static void SubBytes(uint8_t state[4][4]) {
    for (int i = 0; i < 4; ++i)
        for (int j = 0; j < 4; ++j)
            state[i][j] = sbox[state[i][j]];
}

// ShiftRows：行移位操作（循环左移）
static void ShiftRows(uint8_t state[4][4]) {
    uint8_t temp;
    // 第 1 行左移 1
    temp = state[1][0];
    state[1][0] = state[1][1];
    state[1][1] = state[1][2];
    state[1][2] = state[1][3];
    state[1][3] = temp;
    // 第 2 行左移 2
    temp = state[2][0];
    uint8_t temp2 = state[2][1];
    state[2][0] = state[2][2];
    state[2][1] = state[2][3];
    state[2][2] = temp;
    state[2][3] = temp2;
    // 第 3 行左移 3
    temp = state[3][0];
    temp2 = state[3][1];
    uint8_t temp3 = state[3][2];
    state[3][0] = state[3][3];
    state[3][1] = temp;
    state[3][2] = temp2;
    state[3][3] = temp3;
}

// xtime：有限域 GF(2^8) 中乘以 2 的操作（用于 MixColumns）
static uint8_t xtime(uint8_t x) { return (uint8_t)((x << 1) ^ ((x >> 7) & 1 ? 0x1b : 0x00)); }

// MixColumns：对状态的每一列进行线性变换
static void MixColumns(uint8_t state[4][4]) {
    for (int i = 0; i < 4; ++i) {
        uint8_t a = state[0][i];
        uint8_t b = state[1][i];
        uint8_t c = state[2][i];
        uint8_t d = state[3][i];

        uint8_t e = a ^ b ^ c ^ d;
        uint8_t xa = a ^ b; xa = xtime(xa);
        uint8_t xb = b ^ c; xb = xtime(xb);
        uint8_t xc = c ^ d; xc = xtime(xc);
        uint8_t xd = d ^ a; xd = xtime(xd);

        state[0][i] ^= xa ^ e;
        state[1][i] ^= xb ^ e;
        state[2][i] ^= xc ^ e;
        state[3][i] ^= xd ^ e;
    }
}

// Cipher：对单个 16 字节块执行 AES 加密流程（10 轮）
static void Cipher(const uint8_t in[16], uint8_t out[16], const uint8_t RoundKey[176]) {
    uint8_t state[4][4];
    // 把输入按列填充到状态矩阵（列优先）
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            state[j][i] = in[i * 4 + j];
        }
    }

    AddRoundKey(state, RoundKey);

    for (int round = 1; round <= 9; ++round) {
        SubBytes(state);
        ShiftRows(state);
        MixColumns(state);
        AddRoundKey(state, RoundKey + round * 16);
    }

    SubBytes(state);
    ShiftRows(state);
    AddRoundKey(state, RoundKey + 160);

    // 把状态矩阵按列输出为字节序列
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            out[i * 4 + j] = state[j][i];
        }
    }
}

// AES ECB 模式单块加密（用于 CBC 中的块加密）
static void AES_ECB_encrypt(const uint8_t* RoundKey, uint8_t* buf) {
    uint8_t out[16];
    Cipher(buf, out, RoundKey);
    memcpy(buf, out, 16);
}

// 初始化 AES 上下文：做密钥扩展并设置 IV
void AES_init_ctx_iv(struct AES_ctx* ctx, const uint8_t* key, const uint8_t* iv) {
    KeyExpansion(ctx->RoundKey, key);
    if (iv) memcpy(ctx->Iv, iv, 16);
}

// 在 CBC 中，对每个分组先与 IV 异或
static void XorWithIv(uint8_t* buf, const uint8_t* iv) {
    for (uint8_t i = 0; i < 16; ++i) buf[i] ^= iv[i];
}

// AES-CBC 模式的缓冲区加密（长度必须是 16 的倍数）
void AES_CBC_encrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, size_t length) {
    uint8_t* iv = ctx->Iv;
    for (size_t i = 0; i < length; i += 16) {
        XorWithIv(buf + i, iv);
        AES_ECB_encrypt(ctx->RoundKey, buf + i);
        // 更新 IV 为当前密文块（CBC 链接）
        memcpy(iv, buf + i, 16);
    }
}

// 下面实现 CBC 解密所需的逆向操作（逆 S-box、逆行移位、逆列混合等）

// 逆 S 盒，用于解密中的 InvSubBytes
static const uint8_t rsbox[256] = {
    0x52,0x09,0x6A,0xD5,0x30,0x36,0xA5,0x38,0xBF,0x40,0xA3,0x9E,0x81,0xF3,0xD7,0xFB,
    0x7C,0xE3,0x39,0x82,0x9B,0x2F,0xFF,0x87,0x34,0x8E,0x43,0x44,0xC4,0xDE,0xE9,0xCB,
    0x54,0x7B,0x94,0x32,0xA6,0xC2,0x23,0x3D,0xEE,0x4C,0x95,0x0B,0x42,0xFA,0xC3,0x4E,
    0x08,0x2E,0xA1,0x66,0x28,0xD9,0x24,0xB2,0x76,0x5B,0xA2,0x49,0x6D,0x8B,0xD1,0x25,
    0x72,0xF8,0xF6,0x64,0x86,0x68,0x98,0x16,0xD4,0xA4,0x5C,0xCC,0x5D,0x65,0xB6,0x92,
    0x6C,0x70,0x48,0x50,0xFD,0xED,0xB9,0xDA,0x5E,0x15,0x46,0x57,0xA7,0x8D,0x9D,0x84,
    0x90,0xD8,0xAB,0x00,0x8C,0xBC,0xD3,0x0A,0xF7,0xE4,0x58,0x05,0xB8,0xB3,0x45,0x06,
    0xD0,0x2C,0x1E,0x8F,0xCA,0x3F,0x0F,0x02,0xC1,0xAF,0xBD,0x03,0x01,0x13,0x8A,0x6B,
    0x3A,0x91,0x11,0x41,0x4F,0x67,0xDC,0xEA,0x97,0xF2,0xCF,0xCE,0xF0,0xB4,0xE6,0x73,
    0x96,0xAC,0x74,0x22,0xE7,0xAD,0x35,0x85,0xE2,0xF9,0x37,0xE8,0x1C,0x75,0xDF,0x6E,
    0x47,0xF1,0x1A,0x71,0x1D,0x29,0xC5,0x89,0x6F,0xB7,0x62,0x0E,0xAA,0x18,0xBE,0x1B,
    0xFC,0x56,0x3E,0x4B,0xC6,0xD2,0x79,0x20,0x9A,0xDB,0xC0,0xFE,0x78,0xCD,0x5A,0xF4,
    0x1F,0xDD,0xA8,0x33,0x88,0x07,0xC7,0x31,0xB1,0x12,0x10,0x59,0x27,0x80,0xEC,0x5F,
    0x60,0x51,0x7F,0xA9,0x19,0xB5,0x4A,0x0D,0x2D,0xE5,0x7A,0x9F,0x93,0xC9,0x9C,0xEF,
    0xA0,0xE0,0x3B,0x4D,0xAE,0x2A,0xF5,0xB0,0xC8,0xEB,0xBB,0x3C,0x83,0x53,0x99,0x61,
    0x17,0x2B,0x04,0x7E,0xBA,0x77,0xD6,0x26,0xE1,0x69,0x14,0x63,0x55,0x21,0x0C,0x7D
};

// 在 GF(2^8) 上按多项式乘法（用于 InvMixColumns）
static uint8_t Multiply(uint8_t x, uint8_t y) {
    uint8_t result = 0;
    while (y) {
        if (y & 1) result ^= x;
        x = xtime(x);
        y >>= 1;
    }
    return result;
}

// 逆列混合（解密用）
static void InvMixColumns(uint8_t state[4][4]) {
    for (int i = 0; i < 4; ++i) {
        uint8_t a = state[0][i];
        uint8_t b = state[1][i];
        uint8_t c = state[2][i];
        uint8_t d = state[3][i];

        state[0][i] = Multiply(a,0x0e) ^ Multiply(b,0x0b) ^ Multiply(c,0x0d) ^ Multiply(d,0x09);
        state[1][i] = Multiply(a,0x09) ^ Multiply(b,0x0e) ^ Multiply(c,0x0b) ^ Multiply(d,0x0d);
        state[2][i] = Multiply(a,0x0d) ^ Multiply(b,0x09) ^ Multiply(c,0x0e) ^ Multiply(d,0x0b);
        state[3][i] = Multiply(a,0x0b) ^ Multiply(b,0x0d) ^ Multiply(c,0x09) ^ Multiply(d,0x0e);
    }
}

// 逆行移位
static void InvShiftRows(uint8_t state[4][4]) {
    uint8_t temp;
    // 第1行右移1
    temp = state[1][3];
    state[1][3] = state[1][2];
    state[1][2] = state[1][1];
    state[1][1] = state[1][0];
    state[1][0] = temp;

    // 第2行右移2
    temp = state[2][0]; uint8_t temp2 = state[2][1];
    state[2][0] = state[2][2]; state[2][1] = state[2][3]; state[2][2] = temp; state[2][3] = temp2;

    // 第3行右移3
    temp = state[3][0]; temp2 = state[3][1]; uint8_t temp3 = state[3][2];
    state[3][0] = state[3][1]; state[3][1] = state[3][2]; state[3][2] = state[3][3]; state[3][3] = temp;
}

// InvSubBytes：用逆 S 盒将字节替换回原值
static void InvSubBytes(uint8_t state[4][4]) {
    for (int i = 0; i < 4; ++i)
        for (int j = 0; j < 4; ++j)
            state[i][j] = rsbox[state[i][j]];
}

// 逆变换：对单块执行 AES 解密流程
static void InvCipher(const uint8_t in[16], uint8_t out[16], const uint8_t RoundKey[176]) {
    uint8_t state[4][4];
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            state[j][i] = in[i * 4 + j];
        }
    }

    AddRoundKey(state, RoundKey + 160);
    for (int round = 9; round >= 1; --round) {
        InvShiftRows(state);
        InvSubBytes(state);
        AddRoundKey(state, RoundKey + round * 16);
        InvMixColumns(state);
    }
    InvShiftRows(state);
    InvSubBytes(state);
    AddRoundKey(state, RoundKey);

    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            out[i * 4 + j] = state[j][i];
        }
    }
}

// ECB 模式单块解密
static void AES_ECB_decrypt(const uint8_t* RoundKey, uint8_t* buf) {
    uint8_t out[16];
    InvCipher(buf, out, RoundKey);
    memcpy(buf, out, 16);
}

// CBC 模式解密：使用 IV（ctx->Iv）并按块解密，然后与上一密文块异或恢复明文
void AES_CBC_decrypt_buffer(struct AES_ctx* ctx, uint8_t* buf, size_t length) {
    uint8_t prev[16];
    memcpy(prev, ctx->Iv, 16);
    for (size_t i = 0; i < length; i += 16) {
        uint8_t tmp[16];
        memcpy(tmp, buf + i, 16);
        AES_ECB_decrypt(ctx->RoundKey, buf + i);
        for (int j = 0; j < 16; ++j) buf[i + j] ^= prev[j];
        memcpy(prev, tmp, 16);
    }
    // 恢复 ctx->Iv 以便重复使用
    memcpy(ctx->Iv, prev, 16);
}


